package com.jorado.pulsar.client.trace;

import brave.Tracer;
import brave.internal.HexCodec;
import brave.propagation.TraceContext;
import com.jorado.core.utility.SnowFlake;
import com.jorado.core.utility.UniqueIdUtils;
import com.jorado.trace.ZipkinContext;
import com.jorado.trace.util.RequestID;
import org.apache.pulsar.shade.org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class TraceContextContainer {

    private static final Logger LOG = LoggerFactory.getLogger(TraceContextContainer.class);

    private String spanName;

    public TraceContextContainer(String spanName, String appName) {
        this.spanName = spanName;
        synchronized (TraceContextContainer.class) {
            ZipkinContext.startZipkin(1.0f, appName);
        }
    }

    public TraceContext initZipkinContextWithoutScope() {
        long currentSpanId = 0L;
        Boolean sampled = true;
        brave.Span client = getTracer().currentSpan();
        if (client != null) {
            currentSpanId = client.context().spanId();
            sampled = client.context().sampled();
        }
        return initZipkinContext(RequestID.getRequestID(), currentSpanId, sampled, false);
    }

    public TraceContext initZipkinContext(String requestId, Long parentId, Boolean sampled, boolean withScope) {
        long traceIdHigh = 0L, traceIdLow = 0L;
        TraceContext traceCtx = null;
        brave.Span client = null;
        if (StringUtils.isEmpty(requestId)) {
            requestId = UniqueIdUtils.generate();
            RequestID.setRequestID(requestId);
            traceIdHigh = HexCodec.lowerHexToUnsignedLong(requestId.substring(0, 16));
            traceIdLow = HexCodec.lowerHexToUnsignedLong(requestId.substring(16, 32));
            traceCtx = brave.propagation.TraceContext.newBuilder().traceId(traceIdLow).traceIdHigh(traceIdHigh)
                    .spanId(SnowFlake.getInstance().nextId()).parentId(parentId).sampled(sampled).build();
        } else {
            RequestID.setRequestID(requestId);
            if (requestId.contains("-")) {
                LOG.warn("the trace id contain '-' character!");
            } else if (requestId.length() < 32) {
                LOG.warn("the trace id's length is less than 32 characters!");
            } else {
                traceIdHigh = HexCodec.lowerHexToUnsignedLong(requestId.substring(0, 16));
                traceIdLow = HexCodec.lowerHexToUnsignedLong(requestId.substring(16, 32));
            }
            traceCtx = brave.propagation.TraceContext.newBuilder().traceId(traceIdLow).traceIdHigh(traceIdHigh)
                    .spanId(SnowFlake.getInstance().nextId()).parentId(parentId).sampled(sampled).build();

        }
        try {
            client = getTracer().toSpan(traceCtx).name(spanName).start();
            if (withScope) {
                getTracer().withSpanInScope(client);
            }
        } finally {
            if (client != null) {
                client.finish();
            }
        }
        return traceCtx;
    }

    public Tracer getTracer() {
        return ZipkinContext.inst.getTracing().tracer();
    }
}
